import numpy as np

# 飞行模式枚举
class FlightMode:
    WAYPOINT = 1
    RETURN_HOME = 2
    LANDING = 3

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长
SIM_TIME = 30      # 仿真时长（秒）

# 飞行参数
WAYPOINT_RADIUS = 0.1 # 到达航路点的判定半径
HOME_POSITION = np.array([0.0, 0.0, 0.0])  # 返航点位置

# 航路点列表
WAYPOINTS = [
    np.array([2.0, 2.0, 2.0]),
    np.array([2.0, -2.0, 2.0]),
    np.array([-2.0, -2.0, 2.0]),
    np.array([-2.0, 2.0, 2.0])
]

# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 状态变量
pos = np.array([0.0, 0.0, 0.0])   # 位置 x, y, z
vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
current_waypoint_index = 0  # 当前航路点索引
current_mode = FlightMode.WAYPOINT  # 当前飞行模式

# PID控制器初始化
pid_x = PID(2.0, 0.0, 1.0)
pid_y = PID(2.0, 0.0, 1.0)
pid_z = PID(10.0, 2.0, 5.0)

def calculate_velocity_setpoint(target_pos, current_pos):
    """计算速度设定点，考虑最大速度和接近目标时开始减速"""
    direction = target_pos - current_pos
    distance = np.linalg.norm(direction)
    
    if distance < 0.01:  # 防止除以零
        return np.zeros(3)
    
    direction_unit = direction / distance
    MAX_SPEED = 6.0    # 最大飞行速度 (m/s)
    APPROACH_DISTANCE = 1.0  # 接近航路点时开始减速
    
    if distance > APPROACH_DISTANCE:
        return direction_unit * MAX_SPEED
    else:
        speed_factor = distance / APPROACH_DISTANCE
        return direction_unit * MAX_SPEED * speed_factor

def update_flight_mode():
    global current_waypoint_index
    
    current_waypoint = WAYPOINTS[current_waypoint_index]
    distance = np.linalg.norm(pos - current_waypoint)
    
    if distance < WAYPOINT_RADIUS:
        current_waypoint_index += 1
        if current_waypoint_index >= len(WAYPOINTS):
            current_mode = FlightMode.RETURN_HOME
            pid_x.set_setpoint(HOME_POSITION[0])
            pid_y.set_setpoint(HOME_POSITION[1])
            pid_z.set_setpoint(HOME_POSITION[2])
        else:
            next_waypoint = WAYPOINTS[current_waypoint_index]
            pid_x.set_setpoint(next_waypoint[0])
            pid_y.set_setpoint(next_waypoint[1])
            pid_z.set_setpoint(next_waypoint[2])

# 主仿真循环
for frame in range(int(SIM_TIME / DT)):
    # 更新飞行模式
    update_flight_mode()
    
    # 计算速度设定点
    target_pos = WAYPOINTS[current_waypoint_index] if current_mode == FlightMode.WAYPOINT else HOME_POSITION
    velocity_setpoint = calculate_velocity_setpoint(target_pos, pos)
    vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
    
    # 控制器输出
    pid_vx = PID(2.0, 0.0, 0.5)
    pid_vy = PID(2.0, 0.0, 0.5)
    pid_vz = PID(2.0, 0.0, 0.5)

    pid_vx.set_setpoint(vx_setpoint)
    pid_vy.set_setpoint(vy_setpoint)
    pid_vz.set_setpoint(vz_setpoint)
    
    fx = pid_vx.update(vel[0], DT)
    fy = pid_vy.update(vel[1], DT)
    fz = pid_vz.update(vel[2], DT) + MASS * GRAVITY

    force = np.array([fx, fy, fz])
    acc = force / MASS - np.array([0, 0, GRAVITY])
    vel += acc * DT
    pos += vel * DT
    
    if pos[2] < 0:
        pos[2] = 0
        vel[2] = 0

# 完成仿真
